from pycocotools.coco import COCO
import numpy as np
import skimage.io as io
import matplotlib.pyplot as plt
import pylab
import matplotlib; matplotlib.use('TkAgg')
import zipfile
import cv2
# from Polygon import Polygon
from matplotlib.patches import Polygon
from matplotlib.collections import PatchCollection
pylab.rcParams['figure.figsize'] = (8.0, 10.0)

dataDir='D:\\1_DataSets\\coco\\coco2017'
dataType='val2017'
annFile='{}/annotations/instances_{}.json'.format(dataDir,dataType)
# 初始化标注数据的 COCO api
coco=COCO(annFile)

# display COCO categories and supercategories
cats = coco.loadCats(coco.getCatIds())
nms=[cat['name'] for cat in cats]
print('COCO categories: \n{}\n'.format(' '.join(nms)))

nms = set([cat['supercategory'] for cat in cats])
print('COCO supercategories: \n{}'.format(' '.join(nms)))

# get all images containing given categories, select one at random
catIds = coco.getCatIds(catNms=['person','dog','skateboard']);
imgIds = coco.getImgIds(catIds=catIds );
imgIds = coco.getImgIds(imgIds = [324158])
# loadImgs() 返回的是只有一个元素的列表, 使用[0]来访问这个元素
# 列表中的这个元素又是字典类型, 关键字有: ["license", "file_name",
#  "coco_url", "height", "width", "date_captured", "id"]
img = coco.loadImgs(imgIds[np.random.randint(0,len(imgIds))])[0]

# 加载并显示图片,可以使用两种方式: 1) 加载本地图片, 2) 在线加载远程图片
# 1) 使用本地路径, 对应关键字 "file_name"
# I = io.imread('%s/images/%s/%s'%(dataDir,dataType,img['file_name']))

# 2) 使用 url, 对应关键字 "coco_url"
# I = io.imread(img['coco_url'])

# print(img) # a dict

val_z = zipfile.ZipFile(dataDir + "\\" + dataType + ".zip")

def buffer2array(Z, image_name):
    image_buffer = Z.read(image_name)
    image_array = np.frombuffer(image_buffer, dtype='B')
    image_array = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
    return image_array

I = buffer2array(val_z, '%s/%s' % (dataType, img['file_name']))

plt.axis('off')
plt.imshow(I)
plt.show()

# 加载并显示标注信息
plt.imshow(I); plt.axis('off')
annIds = coco.getAnnIds(imgIds=img['id'], catIds=catIds, iscrowd=None)
anns = coco.loadAnns(annIds)
coco.showAnns(anns)
plt.show()

# print(anns)
# print(len(anns))

# # 为 person keypoints 标注信息创建一个 coco 对象
# annFile = '{}/annotations/person_keypoints_{}.json'.format(dataDir,dataType)
# coco_kps=COCO(annFile)

# # load and display keypoints annotations
# plt.imshow(I); plt.axis('off')
# ax = plt.gca()
# annIds = coco_kps.getAnnIds(imgIds=img['id'], catIds=catIds, iscrowd=None)
# anns = coco_kps.loadAnns(annIds)
# coco_kps.showAnns(anns)
# plt.show()


# # 为 caption 标注信息创建一个 coco 对象
# annFile = '{}/annotations/captions_{}.json'.format(dataDir,dataType)
# coco_caps=COCO(annFile)

# # 加载并打印 caption 标注信息
# annIds = coco_caps.getAnnIds(imgIds=img['id']);
# anns = coco_caps.loadAnns(annIds)
# coco_caps.showAnns(anns)
# plt.imshow(I); plt.axis('off'); plt.show()


###-------------------show bbox---------------------------------

def showBBox(self, image, anns, label_box=True):
    """
    show bounding box of annotations or predictions
    anns: loadAnns() annotations or predictions subject to coco results format
    label_box: show background of category labels or not
    """
    if len(anns) == 0:
        return 0
    plt.imshow(image)
    plt.axis('off')
    ax = plt.gca()
    ax.set_autoscale_on(False)
    polygons = []
    color = []
    image2color=dict()
    for cat in self.getCatIds():
        image2color[cat]=(np.random.random((1, 3))*0.7+0.3).tolist()[0]
    for ann in anns:
        c=image2color[ann['category_id']]
        [bbox_x, bbox_y, bbox_w, bbox_h] = ann['bbox']
        poly = [[bbox_x, bbox_y], [bbox_x, bbox_y+bbox_h], [bbox_x+bbox_w, bbox_y+bbox_h], [bbox_x+bbox_w, bbox_y]]
        np_poly = np.array(poly).reshape((4,2))
        polygons.append(Polygon(np_poly))
        color.append(c)
        # option for dash-line
        # ax.add_patch(Polygon(np_poly, linestyle='--', facecolor='none', edgecolor=c, linewidth=2))
        if label_box:
            label_bbox=dict(facecolor=c)
        else:
            label_bbox=None
        if 'score' in ann:
            ax.text(bbox_x, bbox_y, '%s: %.2f'%(self.loadCats(ann['category_id'])[0]['name'], ann['score']), color='white', bbox=label_bbox)
        else:
            ax.text(bbox_x, bbox_y, '%s'%(self.loadCats(ann['category_id'])[0]['name']), color='white', bbox=label_bbox)
    # option for filling bounding box
    # p = PatchCollection(polygons, facecolor=color, linewidths=0, alpha=0.4)
    # ax.add_collection(p)
    p = PatchCollection(polygons, facecolor='none', edgecolors=color, linewidths=2)
    ax.add_collection(p)
    plt.show()

showBBox(coco, I, anns, True)
